#include "ArithmeticEncoder.h"
#include <string.h>
#include <stdio.h>

ArithmeticEncoder::ArithmeticEncoder(ProbabilityModel* model)
{
	this->model = model;
	this->ctwTree = NULL;
	lowerBound = 0;
	upperBound = 0xFFFF;
	underflow = 0;
	symbolSize = model->getSymbolSize();
}

ArithmeticEncoder::ArithmeticEncoder(ContextTree* ctwTree, int modelSize, int blockSize)
{
	this->model = NULL;
	this->ctwTree = ctwTree;
	lowerBound = 0;
	upperBound = 0xFFFF;
	underflow = 0;
	this->modelSize = modelSize;
	this->blockSize = blockSize;
	symbolSize = 1;
}


ArithmeticEncoder::~ArithmeticEncoder(void)
{
}


BitBuffer ArithmeticEncoder::encodeString(const char* string, unsigned long long* symbolsEncoded, int size)
{
	BitBuffer buffer;
	
	if (model != NULL)
	{
		buffer = BitBuffer(size * 8, model->getSymbolSize());
		symbolSize = model->getSymbolSize();
	}
	else
	{
		buffer = BitBuffer(size * 8, 1);
		symbolSize = 1;
	}
	

	unsigned char bitmask = 0x80;
	unsigned char shift = 7;
	*symbolsEncoded = 0;
	unsigned int i = 0;
	int symbolsInTree = 0;
	// init first characters in tree
	if (ctwTree != NULL)
	{
		i += modelSize;
		BitBuffer treeBuffer(modelSize, 1);
		for (int j = 0; j < modelSize; ++j)
		{
			treeBuffer.appendBit((string[j / 8] >> shift) & 0x01);
			buffer.appendBit((string[j / 8] >> shift) & 0x01);
			--shift;
			bitmask >>= 1;
			if (bitmask == 0)
			{
				bitmask = 0x80;
				shift = 7;
			}
		}
		treeBuffer.resetPointers();
		ctwTree->countTree(&treeBuffer, modelSize);
		symbolsInTree = modelSize;
		*symbolsEncoded += modelSize;
	}
	for (; i < size * 8;)
	{
		// tree full
		if (ctwTree != NULL && symbolsInTree % (blockSize * 8) == 0)
		{
			buffer.appendBit((lowerBound >> 14) & 1);
			underflow++;

			while(underflow > 0)
			{
				buffer.appendBit(((lowerBound >> 14) ^ 1) & 1);
				underflow--;
			}
			lowerBound = 0;
			upperBound = 0xFFFF;
			underflow = 0;
			BitBuffer treeBuffer(modelSize, 1);
			for (int j = 0; j < modelSize && i < size * 8; ++j, ++i)
			{
				treeBuffer.appendBit(string[i / 8] >> shift);
				buffer.appendBit(string[i / 8] >> shift);
				--shift;
				bitmask >>= 1;
				if (bitmask == 0)
				{
					bitmask = 0x80;
					shift = 7;
				}
			}
			if (i >= 8 * size)
			{
				return buffer;
			}
			treeBuffer.resetPointers();
			ctwTree->countTree(&treeBuffer, modelSize);
			symbolsInTree = modelSize;
			*symbolsEncoded += modelSize;
		}
		// get symbol from string
		unsigned int symbol = 0;
		for (unsigned int j = 0; j < symbolSize && i < size * 8; ++j, ++i)
		{
			symbol |= (string[i / 8] & bitmask) >> shift;
			--shift;
			bitmask >>= 1;
			if (bitmask == 0)
			{
				bitmask = 0x80;
				shift = 7;
			}
			symbol <<= 1;
		}

		symbol >>= 1;
		//if (symbolSize == 1)
		//{
		//	bacEncodeSymbol(symbol, buffer);
		//}
		//else
		//{
			encodeSymbol(symbol, buffer);
		//}
		
		++symbolsInTree;
		*symbolsEncoded += 1;
	}

	buffer.appendBit((lowerBound >> 14) & 1);
	underflow++;

	while(underflow > 0)
	{
		buffer.appendBit(((lowerBound >> 14) ^ 1) & 1);
		underflow--;
	}

	return buffer;
}

void ArithmeticEncoder::encodeSymbol(unsigned int symbol, BitBuffer& buffer)
{
	int range = upperBound - lowerBound + 1;
	if (symbolSize == 1)
	{
		unsigned int R1;
		if (ctwTree)
		{
			R1 = (unsigned long long)0x10000 * ctwTree->predict(0);
		}
		else
		{
			R1 = (unsigned long long)range * model->getUpperBoundForSymbol(0) / model->getAllSymbolsCount();
		}
		if (symbol == 0)
		{
			if (ctwTree != NULL)
			{
				ctwTree->appendBitToTree(0);
			}
			upperBound = lowerBound + R1 - 1;
		}
		else
		{
			if (ctwTree != NULL)
			{
				R1 = (unsigned long long)0x10000 * ctwTree->appendBitToTree(1);
				lowerBound = upperBound - R1 + 1;
			}
			else
			{
				lowerBound = lowerBound + (range - R1);
			}
		}
	}
	else 
	{
		upperBound = lowerBound + ((unsigned long long)range * model->getUpperBoundForSymbol(symbol)) / model->getAllSymbolsCount() - 1;
		lowerBound = lowerBound + ((unsigned long long)range * model->getLowerBoundForSymbol(symbol)) / model->getAllSymbolsCount();
	}

	while(true)
    {
        if ((upperBound & 0x8000) == (lowerBound & 0x8000))
        {
            buffer.appendBit((upperBound & 0x8000) >> 15);
			while (underflow > 0)
            {
                buffer.appendBit((~upperBound & 0x8000) >> 15);
                --underflow;
            }
        }
        else if ((lowerBound & 0x4000) && !(upperBound & 0x4000))
        {
            ++underflow;
            lowerBound &= 0x3fff;
            upperBound |= 0x4000;
        }
        else
		{
            return;
		}
        lowerBound <<= 1;
        upperBound <<= 1;
        upperBound |= 1;
    }
}

void ArithmeticEncoder::bacEncodeSymbol(unsigned int symbol, BitBuffer& buffer)
{
	int range = upperBound - lowerBound + 1;
	int r1;
	if (ctwTree != NULL)
	{
		r1 = (unsigned long long)0x10000 * ctwTree->predict(0);
	}
	else
	{
		r1 = (unsigned long long)range * model->getUpperBoundForSymbol(0) / model->getAllSymbolsCount();
	}
	
	int r2 = range - r1;
	if (symbol == 0)
	{
		range = r1;
		lowerBound = lowerBound + r2;
		if (ctwTree != NULL)
		{
			ctwTree->appendBitToTree(0);
		}
	}
	else
	{
		range = r2;
		upperBound = lowerBound + r2 - 1;
		if (ctwTree != NULL)
		{
			ctwTree->appendBitToTree(1);
		}
	}

	unsigned int half = 1 << 15;
	unsigned int quarter = 1 << 14;

	while (range <= quarter)
	{
		if (lowerBound >= half)
		{
			buffer.appendBit(1);
			while (underflow > 0)
			{
				buffer.appendBit(0);
				--underflow;
			}
			lowerBound -= half;
			upperBound -= half;
		}
		else if (upperBound <= half)
		{
			buffer.appendBit(0);
			while (underflow > 0)
			{
				buffer.appendBit(1);
				--underflow;
			}
		}
		else
		{
			++underflow;
			lowerBound -= quarter;
			upperBound -= quarter;
		}
		lowerBound <<= 1;
		upperBound <<= 1;
		upperBound |= 1;
		range = upperBound - lowerBound - 1;
	}
}